{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3 variables set to your path：dataroot, saveroot, pair_dir\n",
    "\n",
    "generate npy files，pair files，\n",
    "\n",
    "training dataset demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "from torch.utils.data import Dataset\n",
    "from loadFun import loadmat, kdata2img, multicoilkdata2img\n",
    "from PIL import Image\n",
    "from torchvision import transforms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## generate .npy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def paddingZero_np(np_data: np.array, target_shape: tuple):\n",
    "    H,W = np_data.shape[-2], np_data.shape[-1]\n",
    "    padding_H = target_shape[0] - H\n",
    "    padding_W = target_shape[1] - W\n",
    "    if len(np_data.shape) == 4:\n",
    "        padding_size = ((0,0),(0,0),(padding_H//2, padding_H - padding_H//2), (padding_W//2, padding_W-padding_W//2))\n",
    "    else:\n",
    "        padding_size = ((0,0),(0,0),(0,0),(padding_H//2, padding_H - padding_H//2), (padding_W//2, padding_W-padding_W//2))\n",
    "    padded_np_data = np.pad(np_data, padding_size, mode='constant')\n",
    "    return padded_np_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Task2PngConverter(Dataset):\n",
    "    def __init__(self, root_dir, coil_type, task_type, acc_factor, need_gt=True,saveroot=\"\"):\n",
    "        \"\"\"\n",
    "        root_dir: absolute path of \"ChallengData\"\n",
    "        coil_type: MultiCoil or SingleCoil\n",
    "        task_type: Cine or Mapping\n",
    "        acc_factor: AccFactor04 or AccFactor08 or AccFactor10\n",
    "        \"\"\"\n",
    "        self.name_dict = {\"MultiCoil\":{\"AccFactor04\":\"kspace_sub04\",\n",
    "                                       \"AccFactor08\":\"kspace_sub08\",\n",
    "                                       \"AccFactor10\":\"kspace_sub10\",\n",
    "                                       \"FullSample\":\"kspace_full\"}, \n",
    "                          \"SingleCoil\":{\"AccFactor04\":\"kspace_single_sub04\",\n",
    "                                       \"AccFactor08\":\"kspace_single_sub08\",\n",
    "                                       \"AccFactor10\":\"kspace_single_sub10\",\n",
    "                                       \"FullSample\":\"kspace_single_full\"}}\n",
    "        self.root_dir = root_dir\n",
    "        self.coil_type = coil_type\n",
    "        self.task_type = task_type\n",
    "        self.acc_factor = acc_factor\n",
    "        self.need_gt = need_gt\n",
    "        self.saveroot = saveroot\n",
    "        os.makedirs(self.saveroot, exist_ok=True)\n",
    "\n",
    "        self.train_dir = os.path.join(root_dir, coil_type, task_type, \"TrainingSet\", acc_factor)\n",
    "        self.train_GT_dir = os.path.join(root_dir, coil_type, task_type, \"TrainingSet\", \"FullSample\")\n",
    "        if not self.acc_factor == \"FullSample\":\n",
    "            self.val_dir = os.path.join(root_dir, coil_type, task_type, \"ValidationSet\", acc_factor)\n",
    "            self.val_PNum_list = os.listdir(self.val_dir)\n",
    "\n",
    "        self.train_PNum_list = os.listdir(self.train_dir)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.train_PNum_list)\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        train_item_dir = os.path.join(self.train_dir, self.train_PNum_list[index])\n",
    "        save_root = self.saveroot\n",
    "\n",
    "        train_item_t1_path = os.path.join(train_item_dir, \"T1map.mat\")\n",
    "        train_item_t2_path = os.path.join(train_item_dir, \"T2map.mat\")\n",
    "\n",
    "        # Actually, we won't use the mask data, but we remain the following two lines for some special reasons.\n",
    "        train_item_t1_mask_path = os.path.join(train_item_dir, \"T1map_mask.mat\")\n",
    "        train_item_t2_mask_path = os.path.join(train_item_dir, \"T2map_mask.mat\")\n",
    "        \n",
    "        # GT data\n",
    "        train_itemGT_dir = os.path.join(self.train_GT_dir, self.train_PNum_list[index])\n",
    "        train_itemGT_t1_path = os.path.join(train_itemGT_dir, \"T1map.mat\")\n",
    "        train_itemGT_t2_path = os.path.join(train_itemGT_dir, \"T2map.mat\")\n",
    "\n",
    "        # some patient only has one axis data\n",
    "        if (os.path.exists(train_item_t1_path) & os.path.exists(train_itemGT_t1_path)):\n",
    "            has_t1 = True\n",
    "            self._extract_to_png(train_item_t1_path, train_itemGT_t1_path, save_root, \"t1\")\n",
    "        else:\n",
    "            has_t1 = False\n",
    "        \n",
    "        if (os.path.exists(train_item_t2_path) & os.path.exists(train_itemGT_t2_path)):\n",
    "            has_t2 = True\n",
    "            self._extract_to_png(train_item_t2_path, train_itemGT_t2_path, save_root, \"t2\")\n",
    "        else:\n",
    "            has_t2 = False\n",
    "\n",
    "        return {'t1': train_item_t1_path,\n",
    "                't1_mask': train_item_t1_mask_path,\n",
    "                't1_GT': train_itemGT_t1_path,\n",
    "                't2': train_item_t2_path,\n",
    "                't2_mask': train_item_t2_mask_path,\n",
    "                't2_GT': train_itemGT_t2_path,\n",
    "                'has_t1': has_t1,\n",
    "                'has_t2': has_t2}\n",
    "    \n",
    "    def _kspace2img(self, kspace_raw):\n",
    "        kspace_real = kspace_raw['real']\n",
    "        kspace_imag = kspace_raw['imag']\n",
    "        # kspace_modulus = np.sqrt(kspace_real**2 + kspace_img**2)\n",
    "\n",
    "        kspace_complex = kspace_real + 1j * kspace_imag\n",
    "        iimg = np.fft.ifft2(kspace_complex)\n",
    "        img = np.fft.ifftshift(iimg)\n",
    "        img = np.abs(img)\n",
    "        return img\n",
    "\n",
    "    def _extract_to_png(self, path1: str, path2: str, saveroot: str, mapping_type: str) -> None:\n",
    "        \"\"\"\n",
    "        Extract the .png files from raw .mat file\n",
    "        path1: sampled data path .mat\n",
    "        path2: full data path .mat\n",
    "        dir: to store those extracted .png files.\n",
    "        mapping_type: mark the filename with \"t1\" or \"t2\"\n",
    "        \"\"\"\n",
    "        \n",
    "        patient_id = path1.split(\"/\")[-2]\n",
    "        # 1. process sampled data\n",
    "        mat_data = loadmat(path1)\n",
    "        key1 = self.name_dict[self.coil_type][self.acc_factor]\n",
    "        raw_data = mat_data[key1]\n",
    "        raw_data = paddingZero_np(raw_data, (512,512))\n",
    "        if self.coil_type == \"SingleCoil\":\n",
    "            iimg = kdata2img(raw_data)\n",
    "            frame_num, slice_num, _, _ = iimg.shape\n",
    "            for i in range(frame_num):\n",
    "                for j in range(slice_num):\n",
    "                        img = iimg[i, j, :, :]\n",
    "                        img = img / np.max(img)\n",
    "                        save_path = os.path.join(saveroot, self.coil_type, \"Mapping\", self.acc_factor, patient_id, f\"{self.acc_factor}_{patient_id}_{mapping_type}_{i}_{j}.npy\")\n",
    "                        os.makedirs(os.path.join(saveroot, self.coil_type, \"Mapping\", self.acc_factor, patient_id), exist_ok=True)\n",
    "                        np.save(save_path, img)\n",
    "        elif self.coil_type == \"MultiCoil\":\n",
    "            iimg = multicoilkdata2img(raw_data)\n",
    "            frame_num, slice_num, height, width = iimg.shape\n",
    "            for i in range(frame_num):\n",
    "                for j in range(slice_num):\n",
    "                        img = iimg[i, j, :, :]\n",
    "                        img = img / np.max(img)\n",
    "                        save_path = os.path.join(saveroot, self.coil_type, \"Mapping\", self.acc_factor, patient_id, f\"{self.acc_factor}_{patient_id}_{mapping_type}_{i}_{j}.npy\")\n",
    "                        os.makedirs(os.path.join(saveroot, self.coil_type, \"Mapping\", self.acc_factor, patient_id), exist_ok=True)\n",
    "                        np.save(save_path, img)\n",
    "            pass\n",
    "    \n",
    "    def _preprocess(self):\n",
    "        # print(self.__len__())\n",
    "        for i in range(self.__len__()):\n",
    "            self.__getitem__(i)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "change the dataroot, saveroot in the next cell to your path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataroot = \"/home/txiang/CMRxRecon/CMRxRecon/MICCAIChallenge2023/ChallengeData\"\n",
    "saveroot = \"/home/txiang/CMRxRecon/CMRxRecon/MICCAIChallenge2023/Preprocessed\"\n",
    "for coil in [\"MultiCoil\", \"SingleCoil\"]:\n",
    "    for acc in [\"AccFactor04\", \"AccFactor08\", \"AccFactor10\", \"FullSample\"]:\n",
    "        Task2PngConverter(dataroot,coil,\"Mapping\",acc,saveroot=saveroot)._preprocess()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## generate train pair file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_list = [\"P111\", \"P112\", \"P113\", \"P114\", \"P115\", \"P116\", \"P117\", \"P118\", \"P119\", \"P120\"]\n",
    "class Task2PairFileGenerator():\n",
    "    def __init__(self, root_dir, saveroot, acc):\n",
    "        \"\"\"\n",
    "        root_dir: absolute path of \"ChallengData\"\n",
    "        coil_type: MultiCoil or SingleCoil\n",
    "        task_type: Cine or Mapping\n",
    "        # acc_factor: AccFactor04 or AccFactor08 or AccFactor10\n",
    "        acc: \"04\", \"08\", \"10\"\n",
    "        \"\"\"\n",
    "        self.name_dict = {\"MultiCoil\":{\"AccFactor04\":\"kspace_sub04\",\n",
    "                                       \"AccFactor08\":\"kspace_sub08\",\n",
    "                                       \"AccFactor10\":\"kspace_sub10\",\n",
    "                                       \"FullSample\":\"kspace_full\"}, \n",
    "                          \"SingleCoil\":{\"AccFactor04\":\"kspace_single_sub04\",\n",
    "                                       \"AccFactor08\":\"kspace_single_sub08\",\n",
    "                                       \"AccFactor10\":\"kspace_single_sub10\",\n",
    "                                       \"FullSample\":\"kspace_single_full\"}}\n",
    "        self.root_dir = root_dir\n",
    "        self.coil_type = [\"MultiCoil\", \"SingleCoil\"]\n",
    "        accf = acc\n",
    "        self.acc_factor = [f\"AccFactor{accf}\"]\n",
    "        self.is_train = True\n",
    "        os.makedirs(saveroot, exist_ok=True)\n",
    "        train_pair_file = f\"{saveroot}/Task2_acc_{accf}_train_pair_file_npy.txt\"\n",
    "        val_pair_file = f\"{saveroot}/Task2_acc_{accf}_val_pair_file_npy.txt\"\n",
    "        file_obj_train = open(train_pair_file, \"w\")\n",
    "        file_obj_val = open(val_pair_file, \"w\")\n",
    "\n",
    "        if self.is_train:\n",
    "            for coil_type in self.coil_type:\n",
    "                for acc_factor in self.acc_factor:\n",
    "                    self.train_dir = os.path.join(self.root_dir, coil_type, \"Mapping\", acc_factor)\n",
    "                    self.train_GT_dir = os.path.join(self.root_dir, coil_type, \"Mapping\", \"FullSample\")\n",
    "\n",
    "                    Pnum_list = os.listdir(self.train_dir)\n",
    "                    for Pnum in Pnum_list:\n",
    "                        if Pnum in val_list:\n",
    "                            file_obj = file_obj_val\n",
    "                        else:\n",
    "                            file_obj = file_obj_train\n",
    "                        self.sub_dir = os.path.join(self.train_dir, Pnum)\n",
    "                        self.sub_Gt_dir = os.path.join(self.train_GT_dir, Pnum)\n",
    "                        train_data_list = os.listdir(self.sub_dir)\n",
    "                        for train_data in train_data_list:\n",
    "                            # print(train_data)\n",
    "                            if os.path.exists(os.path.join(self.sub_dir, train_data)) and os.path.exists(os.path.join(self.sub_Gt_dir, train_data.replace(acc_factor, \"FullSample\"))):\n",
    "                                file_obj.writelines([os.path.join(self.sub_dir, train_data), \" \", os.path.join(self.sub_Gt_dir, train_data.replace(acc_factor, \"FullSample\")), \"\\n\"])\n",
    "        file_obj_train.close()\n",
    "        file_obj_val.close()\n",
    "    \n",
    "        file_obj = open(train_pair_file, \"r\")\n",
    "        newfile_name = f\"{saveroot}/Task2_acc_{accf}_train_pair_file_npy_clean.txt\"\n",
    "        newfile_obj = open(newfile_name, \"w\")\n",
    "        lines = file_obj.readlines()\n",
    "        for line in lines:\n",
    "            elem1, elem2 = line.split(\" \")\n",
    "            if not \"FullSample\" in elem1:\n",
    "                newfile_obj.writelines(line)\n",
    "        file_obj.close()\n",
    "        newfile_obj.close()\n",
    "\n",
    "        file_obj = open(val_pair_file, \"r\")\n",
    "        newfile_name = f\"{saveroot}/Task2_acc_{accf}_val_pair_file_npy_clean.txt\"\n",
    "        newfile_obj = open(newfile_name, \"w\")\n",
    "        lines = file_obj.readlines()\n",
    "        for line in lines:\n",
    "            elem1, elem2 = line.split(\" \")\n",
    "            if not \"FullSample\" in elem1:\n",
    "                newfile_obj.writelines(line)\n",
    "        file_obj.close()\n",
    "        newfile_obj.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "change the root_dir, pair_dir in the next cell to your path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "04 ok\n",
      "08 ok\n",
      "10 ok\n"
     ]
    }
   ],
   "source": [
    "# root_dir = \"/home/txiang/CMRxRecon/CMRxRecon/MICCAIChallenge2023/Preprocessed\"\n",
    "root_dir = saveroot\n",
    "pair_dir = \"/home/txiang/CMRxRecon/CMRxRecon_Repo/dataset/train_pair_file\"\n",
    "for acc_ in [\"04\", \"08\", \"10\"]:\n",
    "    Task2PairFileGenerator(root_dir=root_dir, saveroot=pair_dir, acc=acc_)\n",
    "    print(acc_, \"ok\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CMRxReconDataset(Dataset):\n",
    "    \n",
    "    def __init__(self, file_path, transform=None, length=-1):\n",
    "        \"\"\"\n",
    "        root_dir: absolute path of \"ChallengData\"\n",
    "        file_path: the train_pair_file.txt\n",
    "        \"\"\"\n",
    "        self.name_dict = {\"MultiCoil\":{\"AccFactor04\":\"kspace_sub04\",\n",
    "                                       \"AccFactor08\":\"kspace_sub08\",\n",
    "                                       \"AccFactor10\":\"kspace_sub10\",\n",
    "                                       \"FullSample\":\"kspace_full\"}, \n",
    "                          \"SingleCoil\":{\"AccFactor04\":\"kspace_single_sub04\",\n",
    "                                       \"AccFactor08\":\"kspace_single_sub08\",\n",
    "                                       \"AccFactor10\":\"kspace_single_sub10\",\n",
    "                                       \"FullSample\":\"kspace_single_full\"}}\n",
    "        self.file_path = file_path\n",
    "        file_obj = open(self.file_path, \"r\")\n",
    "        self.train_pairs = file_obj.readlines()\n",
    "        if length>0:\n",
    "            self.train_pairs = self.train_pairs[:length]\n",
    "        self.transform = transform\n",
    "        file_obj.close()\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.train_pairs)\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        path, GT_path = self.train_pairs[index].replace(\"\\n\",\"\").split(\" \")\n",
    "        item = np.float32(np.load(path))\n",
    "        GT_item = np.float32(np.load(GT_path))\n",
    "        output = {\"input\": item, \"GT\": GT_item}\n",
    "        if self.transform:\n",
    "            data = np.stack((item, GT_item), axis=-1)\n",
    "            transformed_data = self.transform(data)\n",
    "            output = {\"input\": transformed_data[0,:,:].unsqueeze(0), \n",
    "                      \"GT\": transformed_data[1,:,:].unsqueeze(0), \n",
    "                      \"ipath\":path, \n",
    "                      \"gtpath\":GT_path}\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "15576\n",
      "/home/txiang/CMRxRecon/CMRxRecon/MICCAIChallenge2023/Preprocessed/MultiCoil/Mapping/AccFactor04/P100/AccFactor04_P100_t1_0_3.npy torch.Size([1, 512, 512]) tensor(0.0013) <built-in method max of Tensor object at 0x7fa8a45ca750>\n"
     ]
    }
   ],
   "source": [
    "pair_file = \"/home/txiang/CMRxRecon/CMRxRecon_Repo/dataset/train_pair_file/Task2_acc_04_train_pair_file_npy_clean.txt\"\n",
    "tsfm = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.RandomHorizontalFlip(p=0.5),\n",
    "    transforms.RandomVerticalFlip(p=0.5),\n",
    "])\n",
    "dataset = CMRxReconDataset(pair_file, transform=tsfm, length=-1)\n",
    "print(len(dataset))\n",
    "print(dataset[0][\"ipath\"], dataset[0][\"input\"].shape, dataset[0][\"input\"].min(), dataset[0][\"input\"].max)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "p1",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
